import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone, ServerlessSpec
from collections import Counter
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class PineconeDigitRecognizer:
    def __init__(self, api_key, index_name="mnist-index"):
        self.api_key = api_key
        self.index_name = index_name
        self.pc = None
        self.index = None
        self.initialize_pinecone()
    
    def initialize_pinecone(self):
        """初始化Pinecone客户端并连接到索引"""
        try:
            self.pc = Pinecone(api_key=self.api_key)
            self.index = self.pc.Index(self.index_name)
            logging.info("Pinecone初始化成功，已连接到索引")
            return True
        except Exception as e:
            logging.error(f"Pinecone初始化失败: {e}")
            return False
    
    def preprocess_image(self, image):
        """预处理图像以匹配MNIST数据格式"""
        # 转换为灰度图
        if isinstance(image, np.ndarray):
            if len(image.shape) == 3:
                image = Image.fromarray(image).convert('L')
            else:
                image = Image.fromarray(image)
        
        # 调整大小为8x8像素
        image = image.resize((8, 8))
        
        # 转换为numpy数组并反转颜色（MNIST是白底黑字）
        image_array = np.array(image)
        
        # 如果图像是黑底白字，需要反转
        if np.mean(image_array) < 128:  # 如果整体较暗
            image_array = 255 - image_array
        
        # 归一化到0-16范围（与MNIST数据集一致）
        image_array = (image_array / 255.0) * 16
        
        # 展平为一维向量
        flattened = image_array.ravel()
        
        return flattened.tolist()
    
    def predict_digit(self, image, k=11):
        """使用Pinecone进行数字预测"""
        if image is None:
            return "请先画一个数字！"
        
        try:
            # 预处理图像
            query_vector = self.preprocess_image(image)
            
            # 使用Pinecone进行相似性搜索
            results = self.index.query(
                vector=query_vector,
                top_k=k,
                include_metadata=True
            )
            
            # 提取邻居标签并使用投票机制
            neighbor_labels = [match['metadata']['label'] for match in results['matches']]
            
            if not neighbor_labels:
                return "未找到匹配的数字"
            
            # 投票决定最终预测
            predicted_label = Counter(neighbor_labels).most_common(1)[0][0]
            
            # 获取置信度（最频繁标签的比例）
            confidence = neighbor_labels.count(predicted_label) / len(neighbor_labels)
            
            # 显示详细的邻居信息
            neighbor_info = "\n最近邻匹配: "
            for i, match in enumerate(results['matches'][:5]):  # 显示前5个邻居
                neighbor_info += f"数字{match['metadata']['label']}({match['score']:.3f}) "
            
            return f"识别结果: {predicted_label}\n置信度: {confidence:.2%}{neighbor_info}"
            
        except Exception as e:
            logging.error(f"预测过程中出错: {e}")
            return f"识别错误: {str(e)}"

# 初始化Pinecone识别器
# 你的Pinecone API密钥
API_KEY = "pcsk_5XCTf3_7o7fbVgtePSb8P89EiDJAUNpX2Dc1Dan2GnuML5WQoWTcznFvE5dVhFTmd6A6MC"

recognizer = PineconeDigitRecognizer(API_KEY)

# 定义预测函数
def predict_digit(image):
    """
    预测手写数字
    image: 从画板接收的图像
    """
    return recognizer.predict_digit(image)

# 创建Gradio接口
demo = gr.Interface(
    fn=predict_digit,  # 预测函数
    inputs=gr.Image(
        label="在手写板上画一个数字", 
        image_mode="L", 
        source="canvas", 
        shape=(100, 100),
        type="numpy"
    ),  # 输入：画板
    outputs=gr.Textbox(
        label="识别结果",
        lines=3
    ),  # 输出：文本框显示结果
    title="基于Pinecone云服务的手写数字识别器",
    description="用鼠标在手写板上画一个数字（0-9），AI会通过Pinecone云服务进行识别！第一次推理可能较慢，因为需要初始化云连接。",
    examples=[
        ['example_0.png'],  # 你可以添加示例图片
        ['example_1.png']
    ],
    allow_flagging="never"
)

# 添加CSS样式美化界面
css = """
.container {
    max-width: 800px;
    margin: auto;
}
.output-text {
    font-size: 18px;
    font-weight: bold;
    color: #2E86AB;
}
"""

# 启动Gradio接口
if __name__ == "__main__":
    print("正在启动基于Pinecone的手写数字识别Web应用...")
    print("注意：第一次推理可能较慢，因为需要初始化云连接")
    demo.launch(
        share=True,  # 生成一个公共链接，方便分享
        server_name="0.0.0.0",  # 允许外部访问
        server_port=7860,  # 指定端口
        show_error=True,
        debug=True
    )